package com.sicau.tools;

import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.http.*;
import org.springframework.web.client.RestTemplate;

public class WebSearchTool {

    private final RestTemplate restTemplate;

    // 百度千帆搜索API地址
    private final String API_URL = "https://qianfan.baidubce.com/v2/ai_search/web_search";
    // 替换为有效AccessToken（从百度千帆平台获取）
    private final String ACCESS_TOKEN;

    public WebSearchTool(RestTemplate restTemplate, String accessToken) {
        this.restTemplate = restTemplate;
        this.ACCESS_TOKEN = accessToken;
    }

    /**
     * 调用百度千帆搜索API
     */
    @Tool(description = "百度千帆搜索")
    public String search(@ToolParam(description = "搜索关键词") String keyword) {
        // 构建请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", "Bearer " + ACCESS_TOKEN);

        // 构建过滤条件
        JSONObject filterWeb = JSONUtil.createObj();
        filterWeb.set("type", "web");
        filterWeb.set("top_k", 5);
        JSONObject filterImage = JSONUtil.createObj();
        filterWeb.set("type", "image");
        filterWeb.set("top_k", 5);
        JSONArray filters = JSONUtil.createArray();
        filters.add(filterWeb);
        filters.add(filterImage);

        // 构建请求体（使用Hutool的JSONObject简化JSON构建）
        JSONObject requestBody = JSONUtil.createObj();

        JSONArray messages = JSONUtil.createArray();
        JSONObject message = JSONUtil.createObj()
                .set("role", "user")
                .set("content", keyword);
        messages.add(message);

        requestBody.set("messages", messages);
        requestBody.set("resource_type_filter", filters);

        // 发送POST请求
        HttpEntity<String> requestEntity = new HttpEntity<>(requestBody.toString(), headers);
        ResponseEntity<String> responseEntity = restTemplate.postForEntity(API_URL, requestEntity, String.class);

        if (responseEntity.getStatusCode() != HttpStatus.OK) {
            throw new RuntimeException("API请求失败，状态码：" + responseEntity.getStatusCodeValue());
        }

        try {
            return responseEntity.getBody();
        } catch (Exception e) {
            return "网页抓取失败";
        }
    }
}